/*
 * Copyright (C) 2008 by The Regents of the University of California
 * Redistribution of this file is permitted under the terms of the GNU
 * Public License (GPL).
 *
 * @author Junghoo "John" Cho <cho AT cs.ucla.edu>
 * @date 3/24/2008
 */
 
#include "BTreeIndex.h"
#include "BTreeNode.h"

#define LEAF_CODE -999
#define LEAF_END -999
#define TREE_INFO_PID 0
#define OFFSET_LEAF_NODE 0
#define OFFSET_ROOT_PID 4
#define OFFSET_TREE_HEIGHT 8

using namespace std;

/*
 * BTreeIndex constructor
 */
BTreeIndex::BTreeIndex()
{
		rootPid = 0;
		treeHeight = 0;
		locatePid = rootPid;
		treeOrderIndex = 0;
}

/*
 * Function to change the locatePid.
 * @param pid[IN] the pid to set locatePid
 * @return 0 if change succeeded
 */
RC BTreeIndex::changeLocatePid(PageId pid)
{
	locatePid = pid;
	return 0;
}

/*
 * Open the index file in read or write mode.
 * Under 'w' mode, the index file should be created if it does not exist.
 * @param indexname[IN] the name of the index file
 * @param mode[IN] 'r' for read, 'w' for write
 * @return error code. 0 if no error
 */
RC BTreeIndex::open(const string& indexname, char mode)
{
	openMode = mode;
	
	// Attempt to open file
	if (pf.open(indexname, mode) != 0)
	{
		return RC_FILE_OPEN_FAILED;
	}
	
	// Only need to perform further operation when mode is 'r'
	if (mode == 'r')
	{
		// Attempt to read disk page in memory buffer
		// The first page of index stores information about the tree
		char buffer[PageFile::PAGE_SIZE];
		if (pf.read(TREE_INFO_PID, buffer) != 0)
		{
			return RC_FILE_READ_FAILED;
		}
		
		// Obtain information of node
		memcpy(&rootPid, &(buffer[OFFSET_ROOT_PID]), sizeof(PageId));
		memcpy(&treeHeight, &(buffer[OFFSET_TREE_HEIGHT]), sizeof(int));
		locatePid = rootPid;
	}
	else if (mode == 'w')
	{
		char buffer[PageFile::PAGE_SIZE];
		if (pf.read(TREE_INFO_PID, buffer) == 0)
		{
			// Obtain information of node
			memcpy(&rootPid, &(buffer[OFFSET_ROOT_PID]), sizeof(PageId));
			memcpy(&treeHeight, &(buffer[OFFSET_TREE_HEIGHT]), sizeof(int));
			locatePid = rootPid;
		}
	}
 
	return 0;
}

/*
 * Close the index file.
 * @return error code. 0 if no error
 */
RC BTreeIndex::close()
{
	if (openMode == 'w')
	{
		// Attempt to save information (rootPid, treeHeight)
		// The first page of index stores information about the tree
		char buffer[PageFile::PAGE_SIZE];
		
		memcpy(&(buffer[OFFSET_ROOT_PID]), &rootPid, sizeof(PageId));
		memcpy(&(buffer[OFFSET_TREE_HEIGHT]), &treeHeight, sizeof(int));
		
		if (pf.write(TREE_INFO_PID, buffer) != 0)
		{
			return RC_FILE_WRITE_FAILED;
		}
	}
	
  return pf.close(); // Return the status code of the close
}

/*
 * Function called when insert requires recursion.
 * Every time spliting is needed, this function gets called recursively
 * @param key[IN] the key to insert
 * @param rid[IN] the location of the key to insert
 * @param curPid[IN] the current pid we are trying to insert to
 * @param curLevel[IN] the current node's tree level
 * @param midPid[OUT] the pid of the sibling node
 * @param midKey[OUT] the key in the middle after the split. This key should be inserted to the parent node.
 * @return 0 if successful. Return an error code if there is an error.
 */
RC BTreeIndex::recursiveInsert(int key, const RecordId& rid, PageId curPid, int curLevel,
															 PageId& midPid, int& midKey)
{
	RC retStatus; // To store return code
	
	// Simple case where dealing with leaf node
	if (curLevel == treeHeight)
	{
		// Root node would be a leaf node
		BTLeafNode leafNode;
		leafNode.setParameters(rootPid, treeHeight);
		retStatus = leafNode.read(curPid, pf);
		if (retStatus != 0)
		{
			return retStatus; // RC_FILE_READ_FAILED
		}
		
		// Try the simple insert
		retStatus = leafNode.insert(key, rid);
		
		if (retStatus == 0)
		{
			// Leaf not full, we are done. Write result to disk
			retStatus = leafNode.write(curPid, pf);
			if (retStatus != 0)
			{
				return retStatus; // RC_FILE_WRITE_FAILED
			}
			
			return retStatus; // Insert successful
		}
		else
		{
			// The leaf is currently full, need to insertAndSplit
			
			BTLeafNode siblingLeafNode;
			siblingLeafNode.setParameters(rootPid, treeHeight);
			
			int siblingKey;
			leafNode.insertAndSplit(key, rid, siblingLeafNode, siblingKey);
			
			// The new sibling node would be written to a new pid
			PageId siblingPid = pf.endPid();
			retStatus = siblingLeafNode.write(siblingPid, pf);
			if (retStatus != 0)
			{
				return retStatus; // RC_FILE_WRITE_FAILED
			}
			
			// Set current node's next pointer and write to disk
			leafNode.setNextNodePtr(siblingPid);
			retStatus = leafNode.write(curPid, pf);
			if (retStatus != 0)
			{
				return retStatus; // RC_FILE_WRITE_FAILED
			}
			
			// Need to set output for recursive call
			midPid = siblingPid;
			midKey = siblingKey;
			
			// Special case where we have to create new root
			if (curLevel == 1)
			{
				BTNonLeafNode rootNode;
				retStatus = rootNode.initializeRoot(curPid, siblingKey, siblingPid);
				if (retStatus != 0)
				{
					return retStatus;
				}
				
				rootPid = pf.endPid();
				
				treeHeight++;
				rootNode.setParameters(-1, rootPid, treeHeight);
				
				rootNode.write(rootPid, pf);
				if (retStatus != 0)
				{
					return retStatus; // RC_FILE_WRITE_FAILED
				}
				
				return retStatus; // Success
			}
			else
			{
				// Notify that insertAndSplit was performed
				return RC_NODE_FULL;
			}
		}
	}
	else
	{
		// Case where there is at least one non-leaf node
		
		BTNonLeafNode nonLeafNode;
		nonLeafNode.setParameters(-1, rootPid, treeHeight);
		retStatus = nonLeafNode.read(curPid, pf);
		if (retStatus != 0)
		{
			return retStatus; // RC_FILE_READ_FAILED
		}
		
		// The ret variables are used after recursive call returns
		PageId childPid, retMidPid;
		int retMidKey;
		
		// Find the child pointer to follow
		nonLeafNode.locateChildPtr(key, childPid);
		
		// Recursive call
		retStatus = recursiveInsert(key, rid, childPid, curLevel+1, retMidPid, retMidKey);
		
		// If node to insert to was full
		if (retStatus == RC_NODE_FULL)
		{
			// Insert using the mid values
			RC rc;
			rc = nonLeafNode.insert(retMidKey, retMidPid);
			
			if (rc == 0)
			{
				// Insert was successful, we are done
				retStatus = nonLeafNode.write(curPid, pf);
				if (retStatus != 0)
				{
					return retStatus; // RC_FILE_WRITE_FAILED
				}
				
				return rc; // Insert successful
			}
			else
			{
				// The leaf is currently full, need to insertAndSplit
				
				BTNonLeafNode siblingNonLeafNode;
				siblingNonLeafNode.setParameters(0, rootPid, treeHeight);
				
				// Find the real key of the key to insert
				// This avoids error of setting wrong value as the key value
				char tempBuffer[PageFile::PAGE_SIZE];
				pf.read(retMidPid, tempBuffer);
				int boolLeaf;
				memcpy(&boolLeaf, &(tempBuffer[0]), sizeof(int));
				int realKey;
				if (boolLeaf == LEAF_CODE)
				{
					memcpy(&realKey, &(tempBuffer[20]), sizeof(int));
				}
				else
				{
					memcpy(&realKey, &(tempBuffer[16]), sizeof(int));
				}
				
				int newMidKey;
				nonLeafNode.insertAndSplit(realKey, retMidPid, siblingNonLeafNode, newMidKey);
				
				retStatus = nonLeafNode.write(curPid, pf);
				if (retStatus != 0)
				{
					return retStatus; // RC_FILE_WRITE_FAILED
				}
				
				// The new sibling node would be written to a new pid
				PageId siblingPid = pf.endPid();
				retStatus = siblingNonLeafNode.write(siblingPid, pf);
				if (retStatus != 0)
				{
					return retStatus; // RC_FILE_WRITE_FAILED
				}
				
				// Need to set output for recursive call
				midPid = siblingPid;
				midKey = newMidKey;
				
				// Special case where we have to create new root
				if (curLevel == 1)
				{
					BTNonLeafNode rootNode;
					retStatus = rootNode.initializeRoot(curPid, newMidKey, siblingPid);
					if (retStatus != 0)
					{
						return retStatus;
					}
					
					rootPid = pf.endPid();
					
					rootNode.write(rootPid, pf);
					if (retStatus != 0)
					{
						return retStatus; // RC_FILE_WRITE_FAILED
					}
					
					treeHeight++;
					
					return retStatus;
				}
				
				return rc;
			}
		}
	}
	
	return retStatus;
}

/*
 * Insert (key, RecordId) pair to the index.
 * @param key[IN] the key for the value inserted into the index
 * @param rid[IN] the RecordId for the record being inserted into the index
 * @return error code. 0 if no error
 */
RC BTreeIndex::insert(int key, const RecordId& rid)
{
	// Simple case of Empty tree
	if (treeHeight == 0)
	{
		// Root node would be a leaf node
		BTLeafNode *leafNode = new BTLeafNode();
		
		// Insert the given entry to new node
		leafNode->insert(key, rid);
		
		// Obtain a new pid in pf
		// The first page is reserved for tree information
		rootPid = pf.endPid();
		if (rootPid == TREE_INFO_PID)
		{
			rootPid++;
		}
		
		treeHeight++;
		
		// Set the initial parameters
		leafNode->setParameters(rootPid, treeHeight);
		
		// Write to disk and return the return status of write
		return leafNode->write(rootPid, pf);
	}
	else
	{
		// Tree has height of at least 1, need to use recursion
		
		PageId midPid;
		int midKey;
		return recursiveInsert(key, rid, rootPid, 1, midPid, midKey);
	}
	
	return 0;
}

/*
 * Find the leaf-node index entry whose key value is larger than or 
 * equal to searchKey, and output the location of the entry in IndexCursor.
 * IndexCursor is a "pointer" to a B+tree leaf-node entry consisting of
 * the PageId of the node and the SlotID of the index entry.
 * Note that, for range queries, we need to scan the B+tree leaf nodes.
 * For example, if the query is "key > 1000", we should scan the leaf
 * nodes starting with the key value 1000. For this reason,
 * it is better to return the location of the leaf node entry 
 * for a given searchKey, instead of returning the RecordId
 * associated with the searchKey directly.
 * Once the location of the index entry is identified and returned 
 * from this function, you should call readForward() to retrieve the
 * actual (key, rid) pair from the index.
 * @param key[IN] the key to find.
 * @param cursor[OUT] the cursor pointing to the first index entry
 *                    with the key value.
 * @return error code. 0 if no error.
 */
RC BTreeIndex::locate(int searchKey, IndexCursor& cursor)
{
	// Obtain first 4 bits of buffer to determine whether
	// current page is a leaf node
	char buffer[PageFile::PAGE_SIZE];
	if (pf.read(locatePid, buffer) != 0)
	{
		return RC_FILE_READ_FAILED;
	}
	
	int isLeaf;
	memcpy(&isLeaf, &(buffer[OFFSET_LEAF_NODE]), sizeof(int));
	
	if (treeHeight == 0)
	{
		// B+ tree does not exist
		// Need to create root node
		cursor.pid = 0;
		cursor.eid = 0;
		treeHeight++;
		rootPid = 0;
		locatePid = rootPid;
	}
	else if (isLeaf == LEAF_CODE || treeHeight == 1)
	{
		// Tree consists of only the leaf node, at the root
		BTLeafNode leafNode;
		
		// Read contents to leaf node
		if (leafNode.read(locatePid, pf) != 0)
		{
			return RC_FILE_READ_FAILED; // Error status if read failed
		}
		
		// Locate the eid in the leaf node
		int eid;
		if (leafNode.locate(searchKey, eid) != 0)
		{
			return RC_NO_SUCH_RECORD;
		}
		
		// Return the location of the node
		cursor.pid = locatePid;
		cursor.eid = eid;
		treeOrder[treeOrderIndex] = locatePid;
		locatePid = rootPid;
	}
	else
	{
		// Currently at a non-leaf node, recursively find the leaf node
		BTNonLeafNode nonLeafNode;
		
		// Read contents to non-leaf node
		if (nonLeafNode.read(locatePid, pf) != 0)
		{
			return RC_FILE_READ_FAILED; // Error status if read failed
		}
		
		// Find the child ptr to follow
		int childPid;
		if (nonLeafNode.locateChildPtr(searchKey, childPid) != 0)
		{
			return RC_NO_SUCH_RECORD; // Error status if failed
		}
		
		// Recursively call the locate function
		treeOrder[treeOrderIndex] = locatePid;
		treeOrderIndex++;
		locatePid = childPid;
		IndexCursor locateCursor;
		
		RC retStatus;
		retStatus = locate(searchKey, locateCursor); // Recursive call
		if (retStatus != 0)
		{
			return retStatus;
		}
		
		cursor.pid = locateCursor.pid;
		cursor.eid = locateCursor.eid;
	}
	
  return 0;
}

/*
 * Read the (key, rid) pair at the location specified by the index cursor,
 * and move foward the cursor to the next entry.
 * @param cursor[IN/OUT] the cursor pointing to an leaf-node index entry in the b+tree
 * @param key[OUT] the key stored at the index cursor location.
 * @param rid[OUT] the RecordId stored at the index cursor location.
 * @return error code. 0 if no error
 */
RC BTreeIndex::readForward(IndexCursor& cursor, int& key, RecordId& rid)
{
	BTLeafNode leafNode;
	
	// Read contents to leaf node
	if (leafNode.read(cursor.pid, pf) != 0)
	{
		return RC_FILE_READ_FAILED; // Error status if read failed
	}
	
	// Then read the entry at that index
	if (leafNode.readEntry(cursor.eid, key, rid) != 0)
	{
		return RC_INVALID_RID; // Error status if read failed
	}
	
	// Check if the last entry of the node
	if (cursor.eid < leafNode.getKeyCount() - 1)
	{
		cursor.eid++; // Go to the next index
	}
	else
	{
		int nextNode = leafNode.getNextNodePtr();
		if (nextNode == LEAF_END)
		{
			return RC_END_OF_TREE;
		}
		
		cursor.eid = 0; // Go back to first index
		cursor.pid = nextNode; // Go to the next node
	}	
	
  return 0;
}

/*
 * Purely for testing purposes.
 * Not actually used by the program.
 */
void BTreeIndex::testPrint()
{
	printf("testPrint: %d, %d\n", rootPid, treeHeight);
}
